/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: 集合通信算子信息库
 * Author: jiangchen
 * Create: 2020-6-10
 */

#include "hvd_ops_kernel_builder.h"
#include <securec.h>
#include <functional>
#include "graph/debug/ge_attr_define.h"
#include "ops_kernel_builder_registry.h"
#include "hcom_op_utils.h"
#include "sal_pub.h"


namespace hccl {
REGISTER_OPS_KERNEL_BUILDER(HVD_OPS_LIB_NAME, hccl::HvdOpsKernelBuilder);

HvdOpsKernelBuilder::HvdOpsKernelBuilder()
{
}

HvdOpsKernelBuilder::~HvdOpsKernelBuilder()
{
}

// 返回运行参数，包括workspace 、stream数量以及atomic标志位
ge::Status HvdOpsKernelBuilder::CalcOpRunningParam(ge::Node& node)
{
    HCCL_INFO("calculate hccl runing parameters start.");
    CHK_PRT_RET(!node.GetOpDesc(), HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] GetOpDesc failed. null ptr.", \
        HCOM_ERROR_CODE(HCCL_E_PTR)), ge::INTERNAL_ERROR);

    std::string sCollectiveType = node.GetOpDesc()->GetType();
    HcclResult ret = CheckSupportedOP(sCollectiveType);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] op type[%s] is not supported.",\
        HCOM_ERROR_CODE(ret), sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    // No need to apply sub stream in OP_BASE mode
    if (ge::AttrUtils::SetInt(node.GetOpDesc(), "used_stream_num", 0) == false) {
        HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] op[%s]: to OpDesc failed.", HCOM_ERROR_CODE(HCCL_E_PARA), \
            sCollectiveType.c_str());
        return ge::INTERNAL_ERROR;
    }
    // No need to set workspace mem in OP_BASE mode
    ret = SetOpOutputMemSize(node, sCollectiveType); // set output size
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("errNo[0x%016llx] set op[%s] output size failed.", HCOM_ERROR_CODE(ret),  \
        sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    u64 opMemSize = 0;
    ret = SetOpMemAttr(node, sCollectiveType, opMemSize);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] set op[%s] continuous mem attr failed.",
            HCOM_ERROR_CODE(ret), sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    ret = SetOpAtomicInputIndex(node, sCollectiveType);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] set op[%s] atomic input index failed.",
            HCOM_ERROR_CODE(ret),  sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    HCCL_INFO("calculate hccl runing parameters completed.");
    return ge::SUCCESS;
}

ge::Status HvdOpsKernelBuilder::GenerateTask(const ge::Node &node, ge::RunContext &runContext,
    std::vector<domi::TaskDef> &taskDefList)
{
    HCCL_INFO("GenerateTask start.");
    CHK_PRT_RET(!node.GetOpDesc(), HCCL_ERROR("[Generate][Task]errNo[0x%016llx] opDesc is null.",
        HCOM_ERROR_CODE(HCCL_E_PTR)), ge::INTERNAL_ERROR);

    domi::TaskDef taskDef;
    taskDef.clear_kernel_hccl();
    domi::KernelHcclDef *kernelDefHccl = taskDef.mutable_kernel_hccl();
    CHK_PRT_RET((kernelDefHccl == nullptr), HCCL_ERROR("[Generate][Task]errNo[0x%016llx] kernelDefHccl is null.",
        HCOM_ERROR_CODE(HCCL_E_PTR)), ge::INTERNAL_ERROR);
    taskDef.set_type(RT_MODEL_TASK_HCCL);
    taskDef.set_stream_id(node.GetOpDesc()->GetStreamId());

    std::string sCollectiveType = node.GetOpDesc()->GetType();
    HcclResult ret = CheckSupportedOP(sCollectiveType);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Generate][Task]errNo[0x%016llx] op type[%s] is not supported.",
        HCOM_ERROR_CODE(ret), sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    kernelDefHccl->set_hccl_type(sCollectiveType);
    kernelDefHccl->set_op_index(node.GetOpDesc()->GetId());
    HCCL_DEBUG("GenerateTask: hccl op id[%d].", node.GetOpDesc()->GetId());

    u32 eventID;
    ret = GetEventIDFromOpDesc(node.GetOpDesc(), eventID);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Generate][Task]errNo[0x%016llx] get EventID failed.",
        HCOM_ERROR_CODE(ret)), ge::INTERNAL_ERROR);

    HvdKernelInfoPrivateDef privateDefBuf;
    privateDefBuf.eventID = eventID;
    ret = HcomOpUtils::ConversionOpDataType(node.GetOpDesc(), sCollectiveType, privateDefBuf.dataType);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Generate][Task]errNo[0x%016llx] conversion op data type failed.", HCOM_ERROR_CODE(ret)),
        ge::INTERNAL_ERROR);
    taskDef.set_private_def(&privateDefBuf, sizeof(HvdKernelInfoPrivateDef));
    taskDefList.push_back(taskDef);

    HCCL_INFO("GenerateTask success.");
    return ge::SUCCESS;
}

HcclResult HvdOpsKernelBuilder::GetSupportedOP(std::vector<std::string> &hcclSupportOp) const
{
    hcclSupportOp.assign(HVD_SUPPORTED_OP_TYPE.begin(), HVD_SUPPORTED_OP_TYPE.end());
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelBuilder::SetOpMemAttr(ge::Node &node, const std::string &sCollectiveType, const u64 &opMemSize)
{
    bool bRet =  false;
    // broadcast op use the same buffer for input and output
    bool bRefAttr = (sCollectiveType == HVD_KERNEL_OP_TYPE_BROADCAST);
    bRet = ge::AttrUtils::SetBool(node.GetOpDesc(), ge::ATTR_NAME_REFERENCE, bRefAttr);
    CHK_PRT_RET(!bRet, HCCL_ERROR("[SetOp][MemAttr]errNo[0x%016llx] op[%s]: set  reference attr[%d] to OpDesc"
        "failed.", HCOM_ERROR_CODE(HCCL_E_PARA), sCollectiveType.c_str(), bRefAttr), HCCL_E_PARA);
    if (bRefAttr) {
        bRet = node.GetOpDesc()->UpdateOutputName(node.GetOpDesc()->GetAllInputName());
        CHK_PRT_RET(!bRet, HCCL_ERROR("[SetOp][MemAttr]errNo[0x%016llx] op[%s]: update output name failed.", \
            HCOM_ERROR_CODE(HCCL_E_PARA), sCollectiveType.c_str()), HCCL_E_PARA);
    }
    HCCL_INFO("hvd ops memory attr: op[%s], ref[%d]", sCollectiveType.c_str(), bRefAttr);

    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelBuilder::SetOpAtomicInputIndex(ge::Node& node, const std::string &sCollectiveType)
{
    // allreduce set atomic Input Index attribute
    if (sCollectiveType == HVD_KERNEL_OP_TYPE_ALLREDUCE) {
        vector<int64_t> atomicInputIndex(1, -1); // set value -1 as flag
        if (!ge::AttrUtils::SetListInt(node.GetOpDesc(), "atomic_input_index", atomicInputIndex)) {
            HCCL_ERROR("[SetOp][AtomicInputIndex]errNo[0x%016llx]: set op[%s] atomic index failed.",
                HCOM_ERROR_CODE(HCCL_E_PARA), sCollectiveType.c_str());
            return HCCL_E_PARA;
        }
    }
    return HCCL_SUCCESS;
}

HcclResult HvdOpsKernelBuilder::GetEventIDFromOpDesc(const ge::OpDescPtr &op, u32 &eventID)
{
    std::string event;
    if (ge::AttrUtils::HasAttr(op, "event_id")) {
        if (ge::AttrUtils::GetStr(op, "event_id", event) == false) {
            HCCL_ERROR("[Get][EventID]errNo[0x%016llx]: get event_id failed. no event_id in opDesc",
                HCOM_ERROR_CODE(HCCL_E_PARA));
            return HCCL_E_PARA;
        }
        std::string prefix = "event";
        std::string eventStr = event.substr(prefix.length());
        CHK_RET(SalStrToULong(eventStr, HCCL_BASE_DECIMAL, eventID));
    } else {
        HCCL_ERROR("[Get][EventID]errNo[0x%016llx]: get event_id failed. no event_id in opDesc",
            HCOM_ERROR_CODE(HCCL_E_PARA));
        return HCCL_E_PARA;
    }
    HCCL_DEBUG("get src eventID[%u] success.", eventID);
    return HCCL_SUCCESS;
}
}
